import streamlit as st
import requests
import base64
from PIL import Image
import io
import os

# 页面配置
st.set_page_config(page_title="火山引擎文生图测试", page_icon="🎨", layout="wide")

# 侧边栏配置
st.sidebar.title("API配置")
model_endpoint = st.sidebar.text_input(
    "Model Endpoint",
    placeholder="输入模型API端点URL",
    value="https://ark.cn-beijing.volces.com/api/v3/images/generations",
)

# 在侧边栏添加使用说明
with st.sidebar.expander("ℹ️ 使用说明"):
    st.markdown(
        """
        ### 如何使用
        1. 配置环境变量 `ARK_API_KEY` 为您的火山引擎API密钥（推荐方式）
        2. 或者在页面上方的输入框中手动输入API密钥
        3. 在左侧边栏确认模型API端点URL（默认已填写）
        4. 在上方输入框中输入用于生成图像的提示词
        5. 调整图像生成参数（可选）
        6. 点击"生成图像"按钮
        
        ### 参数说明
        - **提示词**: 描述你想要生成的图像内容
        - **反向提示词**: 指定不希望在图像中出现的内容
        - **宽度/高度**: 生成图像的尺寸（将自动转换为1K/2K/4K）
        - **随机种子**: 控制生成的随机性，相同种子会产生相似结果
        
        ### 技术说明
        - 本工具使用火山引擎的REST API进行图像生成
        - 图像尺寸会自动转换：
          - 小于1024px: 1K
          - 1024px-2047px: 2K
          - 2048px及以上: 4K
        - 生成的图像将以URL形式返回并直接显示
        
        > 注意：环境变量方式更安全，推荐在生产环境中使用
        """
    )

# 初始化session state
if "api_key" not in st.session_state:
    st.session_state.api_key = ""

# 主页面
st.title("🎨 火山引擎文生图API测试")
st.markdown("通过此工具测试火山引擎的文生图模型API")

# 显示模型名称
st.markdown(
    """
    <div style="background-color: #f0f8ff; padding: 10px; border-radius: 8px; border-left: 4px solid #4682b4; margin-bottom: 20px;">
        <p style="margin: 0; font-size: 16px; color: #333;">
            当前模型: <strong style="color: #4682b4;">Doubao-Seedream-4.0</strong> 
            <span style="font-size: 14px; color: #666;">(ep-20250912120004-dwhj4)</span>
        </p>
    </div>
    """,
    unsafe_allow_html=True
)

# 输入区域
prompt = st.text_area(
    "输入提示词", height=100, placeholder="请输入用于生成图像的文本描述..."
)
negative_prompt = st.text_area(
    "反向提示词（可选）", height=100, placeholder="不希望在图像中出现的内容..."
)

# 参数设置
col1, col2, col3 = st.columns(3)
with col1:
    width = st.number_input("宽度", min_value=64, max_value=2048, value=256, step=64)
    steps = st.number_input("迭代步数", min_value=1, max_value=100, value=20)
with col2:
    height = st.number_input("高度", min_value=64, max_value=2048, value=512, step=64)
    guidance_scale = st.slider(
        "指导权重", min_value=1.0, max_value=20.0, value=7.5, step=0.5
    )
with col3:
    num_images = st.number_input("生成数量", min_value=1, max_value=4, value=1)
    seed = st.number_input(
        "随机种子", min_value=-1, value=-1, help="设为-1表示随机种子"
    )

# 生成按钮
generate_btn = st.button("🚀 生成图像", type="primary", use_container_width=True)

# 在主页面添加API密钥输入框作为兜底方案
api_key_from_input = st.text_input(
    "API Key (环境变量不可用时使用)",
    type="password",
    value=st.session_state.api_key,
    placeholder="请输入您的火山引擎API密钥",
    help="如果已配置环境变量ARK_API_KEY，此处无需填写",
)

# 更新session state
if api_key_from_input != st.session_state.api_key:
    st.session_state.api_key = api_key_from_input

# 结果显示区域
if generate_btn:
    # 获取环境变量中的API密钥
    api_key = os.environ.get("ARK_API_KEY")

    # 如果环境变量中没有API密钥，则使用用户输入的API密钥
    if not api_key:
        api_key = api_key_from_input

    if not api_key:
        st.error("请配置环境变量 ARK_API_KEY 或在上方输入API密钥")
    elif not model_endpoint:
        st.error("请在侧边栏填写完整的API配置信息")
    elif not prompt:
        st.error("请输入提示词")
    else:
        with st.spinner("正在生成图像..."):
            try:
                # 构造请求头
                headers = {
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {api_key}",
                }

                # 根据宽度和高度确定尺寸参数
                if width >= 2048 or height >= 2048:
                    size = "4K"
                elif width >= 1024 or height >= 1024:
                    size = "2K"
                else:
                    size = "1K"

                # 构造请求体
                payload = {
                    "model": "ep-20250912120004-dwhj4",  # Doubao-Seedream-4.0
                    "prompt": prompt,
                    "size": size,
                    "response_format": "url",
                    "watermark": False,
                }

                # 添加可选参数
                if negative_prompt:
                    payload["negative_prompt"] = negative_prompt

                if seed != -1:
                    payload["seed"] = seed

                # 发送请求到火山引擎API
                response = requests.post(model_endpoint, headers=headers, json=payload)

                try:
                    result = response.json()
                except Exception as e:
                    st.error(f"无法解析API响应为JSON格式: {str(e)}")
                    st.write(f"响应内容: {response.text}")
                    continue_processing = False
                else:
                    continue_processing = True

                # 只有在JSON解析成功时才继续处理
                if continue_processing:
                    if response.status_code == 200:
                        # 检查API响应格式是否正确
                        if "data" in result and isinstance(result["data"], list):
                            st.subheader("生成结果")
                            images = result["data"]
                            
                            # 创建居中显示的容器
                            for i, img_data in enumerate(images):
                                # 为每个图像创建一个独立的容器
                                with st.container():
                                    st.markdown(f"<h4 style='text-align: center;'>图像 {i+1}</h4>", unsafe_allow_html=True)
                                    if "url" in img_data:
                                        # 使用自定义CSS确保图像完整显示在中央，并在两侧留出空白
                                        # 添加显示框以区分背景
                                        st.markdown(
                                            f"""
                                            <div style="display: flex; justify-content: center; align-items: center; margin: 20px 0;">
                                                <div style="width: 100%; max-width: 800px; padding: 0 20px;">
                                                    <div style="border: 2px solid #e0e0e0; border-radius: 12px; padding: 20px; background-color: #f9f9f9; box-shadow: 0 4px 12px rgba(0,0,0,0.05);">
                                                        <img src="{img_data['url']}" style="max-width: 100%; height: auto; border-radius: 8px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); display: block; margin: 0 auto;" />
                                                    </div>
                                                </div>
                                            </div>
                                            """,
                                            unsafe_allow_html=True
                                        )
                                    else:
                                        st.warning(f"图像 {i+1} 的URL不可用")
                            st.success("图像生成完成！")
                        else:
                            st.error("API返回数据格式不正确")
                            # 显示响应内容以便调试
                            st.write("响应中缺少 'data' 字段或格式不正确")
                            st.write("完整响应内容:", result)
                    else:
                        st.error(f"API调用失败: {response.status_code}")
                        # 直接显示完整的错误响应
                        st.write("错误响应内容:", result)

            except Exception as e:
                st.error(f"图像生成失败: {str(e)}")
